package main

import (
	"fmt"
	"log"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/duke-git/lancet/v2/fileutil"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	_ "github.com/joho/godotenv/autoload"
)

var (
	path     string
	ignores  = map[string]struct{}{}
	only     = map[string]struct{}{}
	mysqlDSN string

	wg sync.WaitGroup

	taskCh = make(chan Task)
)

type Task struct {
	tbname string
}

func initSet() {
	path = os.Getenv("IMPORT_PATH")
	mysqlDSN = os.Getenv("MYSQL_DSN")
	ignoresstr := strings.TrimSpace(os.Getenv("IMPORT_IGNORES"))
	if ignoresstr != "" {
		for _, ignore := range strings.Split(ignoresstr, ",") {
			ignores[ignore] = struct{}{}
		}
	}
	onlystr := strings.TrimSpace(os.Getenv("IMPORT_ONLY"))
	if onlystr != "" {
		for _, o := range strings.Split(onlystr, ",") {
			only[o] = struct{}{}
		}
	}
	logfile, err := os.OpenFile("imt.log", os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		log.Fatal(err)
	}
	log.SetOutput(logfile)
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	numCpu := runtime.NumCPU()
	initSet()
	for i := 0; i < numCpu; i++ {
		go worker(i)
	}
	subs, err := os.ReadDir(path)
	if err != nil {
		log.Fatalln(err)
		return
	}

	for _, sub := range subs {
		if !sub.IsDir() {
			continue
		}
		wg.Add(1)
		taskCh <- Task{sub.Name()}

	}
	wg.Wait()
	log.Println("all done")
}

func execSqlFile(sqlFile string, db *gorm.DB) {
	data, err := fileutil.ReadFileToString(sqlFile)
	if err != nil {
		log.Println(err)
		return
	}

	sqls := strings.Split(data, ";\n")

	for _, sql := range sqls {
		sql = strings.TrimSpace(sql)
		if sql == "" {
			continue
		}
		// log.Println(sql)
		if err := db.Exec(sql).Error; err != nil {
			log.Println(err)
		}
	}

}

func execData(path string, db *gorm.DB) {

	files, err := os.ReadDir(path)
	if err != nil {
		log.Println(err)
		return
	}

	for _, f := range files {
		if !f.IsDir() {
			sqlFile := f.Name()
			if strings.HasSuffix(sqlFile, ".sql") {
				execSqlFile(fmt.Sprintf("%s/%s", path, sqlFile), db)
			}
		}
	}

}

func worker(i int) {
	dd, err := gorm.Open(mysql.Open(mysqlDSN), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Error),
	})
	if err != nil {
		log.Fatalln(err)
		return
	}

	for t := range taskCh {
		log.Printf("worker %d: %s\n", i, t.tbname)
		structure := fmt.Sprintf("%s/%s/structure.sql", path, t.tbname)
		dataDir := fmt.Sprintf("%s/%s/data", path, t.tbname)
		if len(only) > 0 {
			_, ok := only[t.tbname]
			if !ok {
				wg.Done()
				continue
			}
		} else {
			_, ok := ignores[t.tbname]
			if ok {
				wg.Done()
				continue
			}
		}
		if err := dd.Exec("drop table if exists `" + t.tbname + "`").Error; err != nil {
			log.Println(err)
		}
		if fileutil.IsExist(structure) {
			execSqlFile(structure, dd)
		}

		if fileutil.IsExist(dataDir) {
			execData(dataDir, dd)
		}
		log.Println(t.tbname, "done")
		wg.Done()
	}
}
